"""
工具函数模块
"""

import os
import sys
import logging
import torch
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, Any, Optional


def setup_logging(log_file: str, level: int = logging.INFO):
    """设置日志记录
    
    参数:
        log_file: 日志文件路径
        level: 日志级别
    """
    # 创建日志目录
    os.makedirs(os.path.dirname(log_file), exist_ok=True)
    
    # 配置日志格式
    logging.basicConfig(
        level=level,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler(sys.stdout)
        ]
    )


def save_checkpoint(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    scheduler: Optional[torch.optim.lr_scheduler._LRScheduler],
    epoch: int,
    best_val_loss: float,
    path: str
):
    """保存模型检查点
    
    参数:
        model: 模型
        optimizer: 优化器
        scheduler: 学习率调度器
        epoch: 当前轮数
        best_val_loss: 最佳验证损失
        path: 保存路径
    """
    os.makedirs(os.path.dirname(path), exist_ok=True)
    
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'best_val_loss': best_val_loss
    }
    
    if scheduler is not None:
        checkpoint['scheduler_state_dict'] = scheduler.state_dict()
    
    torch.save(checkpoint, path)


def load_checkpoint(
    path: str,
    device: torch.device
) -> Dict[str, Any]:
    """加载模型检查点
    
    参数:
        path: 检查点路径
        device: 计算设备
        
    返回:
        检查点字典
    """
    checkpoint = torch.load(path, map_location=device)
    return checkpoint


def count_parameters(model: torch.nn.Module) -> tuple:
    """计算模型参数数量
    
    参数:
        model: 模型
        
    返回:
        (总参数数, 可训练参数数)
    """
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    return total_params, trainable_params


def set_seed(seed: int):
    """设置随机种子
    
    参数:
        seed: 随机种子
    """
    torch.manual_seed(seed)
    np.random.seed(seed)
    
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        
        # 确保结果可重现
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def get_device(device: str = 'auto') -> torch.device:
    """获取计算设备
    
    参数:
        device: 设备类型 ('auto', 'cpu', 'cuda')
        
    返回:
        PyTorch设备对象
    """
    if device == 'auto':
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    return torch.device(device)


def print_model_info(model: torch.nn.Module, input_shape: tuple):
    """打印模型信息
    
    参数:
        model: 模型
        input_shape: 输入形状
    """
    total_params, trainable_params = count_parameters(model)
    
    print(f"模型参数总数: {total_params:,}")
    print(f"可训练参数数: {trainable_params:,}")
    print(f"输入形状: {input_shape}")
    
    # 计算模型大小
    param_size = 0
    buffer_size = 0
    
    for param in model.parameters():
        param_size += param.nelement() * param.element_size()
    
    for buffer in model.buffers():
        buffer_size += buffer.nelement() * buffer.element_size()
    
    model_size = (param_size + buffer_size) / 1024 / 1024  # MB
    print(f"模型大小: {model_size:.2f} MB")


def plot_loss_curves(
    train_losses: list,
    val_losses: Optional[list] = None,
    save_path: Optional[str] = None,
    title: str = "Loss Curves"
):
    """绘制损失曲线
    
    参数:
        train_losses: 训练损失列表
        val_losses: 验证损失列表
        save_path: 保存路径
        title: 图标题
    """
    plt.figure(figsize=(10, 6))
    
    epochs = range(1, len(train_losses) + 1)
    plt.plot(epochs, train_losses, 'b-', label='Training Loss')
    
    if val_losses is not None:
        plt.plot(epochs, val_losses, 'r-', label='Validation Loss')
    
    plt.title(title)
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.grid(True)
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show()


def plot_learning_curve(
    metrics: Dict[str, list],
    save_path: Optional[str] = None,
    title: str = "Learning Curve"
):
    """绘制学习曲线
    
    参数:
        metrics: 指标字典，键为指标名称，值为指标值列表
        save_path: 保存路径
        title: 图标题
    """
    plt.figure(figsize=(12, 8))
    
    epochs = range(1, len(list(metrics.values())[0]) + 1)
    
    for i, (name, values) in enumerate(metrics.items()):
        plt.subplot(len(metrics), 1, i + 1)
        plt.plot(epochs, values)
        plt.title(name)
        plt.xlabel('Epoch')
        plt.ylabel(name)
        plt.grid(True)
    
    plt.suptitle(title)
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show()


def compare_fields(
    true_field: torch.Tensor,
    pred_field: torch.Tensor,
    field_names: list,
    save_path: Optional[str] = None,
    title: str = "Field Comparison"
):
    """比较真实场和预测场
    
    参数:
        true_field: 真实场，形状为 [nt, nx, ny, nz, n_components]
        pred_field: 预测场，形状与真实场相同
        field_names: 场分量名称列表
        save_path: 保存路径
        title: 图标题
    """
    nt, nx, ny, nz, n_comp = true_field.shape
    
    # 选择中间时间步和z切片
    t_idx = nt // 2
    z_idx = nz // 2
    
    # 创建子图
    fig, axes = plt.subplots(2, n_comp, figsize=(4 * n_comp, 8))
    
    if n_comp == 1:
        axes = axes.reshape(2, 1)
    
    # 绘制每个分量的真实值和预测值
    for i in range(n_comp):
        # 真实值
        im = axes[0, i].imshow(
            true_field[t_idx, :, :, z_idx, i].cpu().numpy(),
            origin='lower',
            cmap='viridis'
        )
        axes[0, i].set_title(f'True {field_names[i]}')
        plt.colorbar(im, ax=axes[0, i])
        
        # 预测值
        im = axes[1, i].imshow(
            pred_field[t_idx, :, :, z_idx, i].cpu().numpy(),
            origin='lower',
            cmap='viridis'
        )
        axes[1, i].set_title(f'Pred {field_names[i]}')
        plt.colorbar(im, ax=axes[1, i])
    
    plt.suptitle(title)
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show()


def plot_error_distribution(
    true_field: torch.Tensor,
    pred_field: torch.Tensor,
    field_names: list,
    save_path: Optional[str] = None,
    title: str = "Error Distribution"
):
    """绘制误差分布
    
    参数:
        true_field: 真实场
        pred_field: 预测场
        field_names: 场分量名称列表
        save_path: 保存路径
        title: 图标题
    """
    nt, nx, ny, nz, n_comp = true_field.shape
    
    # 计算误差
    error = torch.abs(true_field - pred_field)
    rel_error = error / (torch.abs(true_field) + 1e-8)
    
    # 创建子图
    fig, axes = plt.subplots(2, n_comp, figsize=(4 * n_comp, 8))
    
    if n_comp == 1:
        axes = axes.reshape(2, 1)
    
    # 绘制每个分量的绝对误差和相对误差
    for i in range(n_comp):
        # 绝对误差
        abs_err = error[..., i].flatten().cpu().numpy()
        axes[0, i].hist(abs_err, bins=50, density=True)
        axes[0, i].set_title(f'Abs Error {field_names[i]}')
        axes[0, i].set_xlabel('Error')
        axes[0, i].set_ylabel('Density')
        axes[0, i].grid(True)
        
        # 相对误差
        rel_err = rel_error[..., i].flatten().cpu().numpy()
        axes[1, i].hist(rel_err, bins=50, density=True)
        axes[1, i].set_title(f'Rel Error {field_names[i]}')
        axes[1, i].set_xlabel('Relative Error')
        axes[1, i].set_ylabel('Density')
        axes[1, i].grid(True)
    
    plt.suptitle(title)
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show()


def compute_field_metrics(
    true_field: torch.Tensor,
    pred_field: torch.Tensor
) -> Dict[str, float]:
    """计算场指标
    
    参数:
        true_field: 真实场
        pred_field: 预测场
        
    返回:
        指标字典
    """
    # 计算误差
    error = true_field - pred_field
    abs_error = torch.abs(error)
    rel_error = abs_error / (torch.abs(true_field) + 1e-8)
    
    # 计算指标
    mse = torch.mean(error ** 2).item()
    mae = torch.mean(abs_error).item()
    rmse = torch.sqrt(mse)
    
    # 计算相对误差指标
    mean_rel_error = torch.mean(rel_error).item()
    max_rel_error = torch.max(rel_error).item()
    
    # 计算R²分数
    ss_res = torch.sum(error ** 2).item()
    ss_tot = torch.sum((true_field - torch.mean(true_field)) ** 2).item()
    r2 = 1 - (ss_res / ss_tot) if ss_tot != 0 else float('nan')
    
    return {
        'mse': mse,
        'mae': mae,
        'rmse': rmse,
        'mean_rel_error': mean_rel_error,
        'max_rel_error': max_rel_error,
        'r2': r2
    }


def save_metrics(
    metrics: Dict[str, float],
    save_path: str
):
    """保存指标到文件
    
    参数:
        metrics: 指标字典
        save_path: 保存路径
    """
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    
    with open(save_path, 'w') as f:
        for name, value in metrics.items():
            f.write(f"{name}: {value}\n")


def load_metrics(
    load_path: str
) -> Dict[str, float]:
    """从文件加载指标
    
    参数:
        load_path: 加载路径
        
    返回:
        指标字典
    """
    metrics = {}
    
    with open(load_path, 'r') as f:
        for line in f:
            if ':' in line:
                name, value = line.strip().split(':', 1)
                try:
                    metrics[name] = float(value)
                except ValueError:
                    metrics[name] = value
    
    return metrics